from qutip import *
import numpy as np
import matplotlib.pyplot as plt
sys.path.append('D:\Dropbox\Dokumente\PI3\Code\Python\quantum-tools')
import qutip_enhanced
bloch = Bloch3d()
qte = qutip_enhanced.QutipEnhanced()

rho = (fock_dm(2,0 )+fock_dm(2,1)).unit()
h = 0*rho
gamma = 1/5850.
# decay_operator = np.sqrt(gamma)*0.5*(jmat(0.5, 'x') - 1j*jmat(0.5, 'y')) #when decay goes to ground state
decay_operator = np.sqrt(gamma)* (jmat(0.5, 'x') + jmat(0.5, 'y')) #found by trying.. it just gives the right decay
ode = mesolve(h, rho,  np.linspace(0, 20e3, 5), [decay_operator], [jmat(0.5, 'z')])

import matplotlib.pyplot as plt
plt.figure()
plt.plot(ode.times, ode.expect[0])
plt.plot(ode.times, 0.5*np.exp(-ode.times*gamma) + 0)
plt.legend()
plt.show()